import os
import requests
import zipfile
import shutil
from tqdm import tqdm
import random
import cv2
import numpy as np
from PIL import Image, ImageEnhance

# 配置参数
DATA_DIR = 'data'
FRUITS_CLASSES = ['apple', 'banana', 'orange', 'strawberry', 'grape', 'kiwi', 'lemon', 'mango', 'pineapple', 'watermelon']
TRAIN_RATIO = 0.7
VAL_RATIO = 0.2
TEST_RATIO = 0.1

def download_dataset():
    """下载水果数据集（示例使用Kaggle水果数据集）"""
    print("正在下载水果数据集...")
    
    # 创建数据目录
    os.makedirs(DATA_DIR, exist_ok=True)
    
    # 这里应该使用实际的数据集下载链接
    # 由于我们无法直接下载Kaggle数据集，这里只是示例
    # 实际使用时，可以手动下载数据集并放入data目录
    
    print("请手动下载水果数据集并解压到data目录")
    print("推荐数据集: Fruits-360 (https://www.kaggle.com/datasets/moltean/fruits)")
    
    # 创建必要的目录结构
    for split in ['train', 'valid', 'test']:
        for subdir in ['images', 'labels']:
            os.makedirs(os.path.join(DATA_DIR, split, subdir), exist_ok=True)

def create_yolo_dataset():
    """创建YOLO格式的数据集"""
    print("正在创建YOLO格式的数据集...")
    
    # 假设已经有了水果图片，存放在data/fruits目录下
    # 每个水果类别有一个子目录
    fruits_dir = os.path.join(DATA_DIR, 'fruits')
    
    if not os.path.exists(fruits_dir):
        print(f"错误: {fruits_dir} 目录不存在")
        print("请确保已下载并解压水果数据集")
        return
    
    # 获取所有水果类别
    fruit_classes = [d for d in os.listdir(fruits_dir) if os.path.isdir(os.path.join(fruits_dir, d))]
    
    # 创建类别到索引的映射
    class_to_idx = {cls: idx for idx, cls in enumerate(FRUITS_CLASSES)}
    
    # 处理每个水果类别
    for fruit_class in fruit_classes:
        if fruit_class not in FRUITS_CLASSES:
            continue
            
        class_idx = class_to_idx[fruit_class]
        class_dir = os.path.join(fruits_dir, fruit_class)
        
        # 获取该类别的所有图片
        images = [f for f in os.listdir(class_dir) if f.endswith(('.jpg', '.jpeg', '.png'))]
        
        # 随机打乱
        random.shuffle(images)
        
        # 划分数据集
        n_images = len(images)
        n_train = int(n_images * TRAIN_RATIO)
        n_val = int(n_images * VAL_RATIO)
        
        train_images = images[:n_train]
        val_images = images[n_train:n_train+n_val]
        test_images = images[n_train+n_val:]
        
        # 处理训练集
        process_images(class_dir, train_images, 'train', class_idx, fruit_class)
        
        # 处理验证集
        process_images(class_dir, val_images, 'valid', class_idx, fruit_class)
        
        # 处理测试集
        process_images(class_dir, test_images, 'test', class_idx, fruit_class)

def process_images(class_dir, images, split, class_idx, class_name):
    """处理图片并创建YOLO格式的标签"""
    for img_file in tqdm(images, desc=f"处理 {class_name} - {split}"):
        # 源图片路径
        src_path = os.path.join(class_dir, img_file)
        
        # 目标图片路径
        dst_img_path = os.path.join(DATA_DIR, split, 'images', f"{class_name}_{img_file}")
        
        # 目标标签路径
        dst_label_path = os.path.join(DATA_DIR, split, 'labels', f"{class_name}_{img_file.rsplit('.', 1)[0]}.txt")
        
        # 复制图片
        shutil.copy(src_path, dst_img_path)
        
        # 创建YOLO格式的标签
        img = cv2.imread(src_path)
        h, w, _ = img.shape
        
        # 假设物体占据图片的中心区域，约占图片的60-80%
        x_center = 0.5
        y_center = 0.5
        width = random.uniform(0.6, 0.8)
        height = random.uniform(0.6, 0.8)
        
        # 写入标签文件
        with open(dst_label_path, 'w') as f:
            f.write(f"{class_idx} {x_center} {y_center} {width} {height}\n")

def augment_data():
    """数据增强"""
    print("正在进行数据增强...")
    
    train_img_dir = os.path.join(DATA_DIR, 'train', 'images')
    train_label_dir = os.path.join(DATA_DIR, 'train', 'labels')
    
    # 获取所有训练图片
    train_images = [f for f in os.listdir(train_img_dir) if f.endswith(('.jpg', '.jpeg', '.png'))]
    
    for img_file in tqdm(train_images, desc="数据增强"):
        # 图片路径
        img_path = os.path.join(train_img_dir, img_file)
        
        # 标签路径
        label_path = os.path.join(train_label_dir, f"{img_file.rsplit('.', 1)[0]}.txt")
        
        # 读取图片
        img = Image.open(img_path)
        
        # 读取标签
        with open(label_path, 'r') as f:
            label_content = f.read()
        
        # 水平翻转
        flipped_img = img.transpose(Image.FLIP_LEFT_RIGHT)
        flipped_img_path = os.path.join(train_img_dir, f"flip_{img_file}")
        flipped_img.save(flipped_img_path)
        
        # 修改标签（水平翻转需要修改x坐标）
        flipped_label_path = os.path.join(train_label_dir, f"flip_{img_file.rsplit('.', 1)[0]}.txt")
        with open(flipped_label_path, 'w') as f:
            for line in label_content.strip().split('\n'):
                parts = line.split()
                cls_id = parts[0]
                x_center = 1.0 - float(parts[1])  # 水平翻转，x坐标变为1-x
                y_center = parts[2]
                width = parts[3]
                height = parts[4]
                f.write(f"{cls_id} {x_center} {y_center} {width} {height}\n")
        
        # 亮度增强
        enhancer = ImageEnhance.Brightness(img)
        bright_img = enhancer.enhance(1.5)  # 增加亮度
        bright_img_path = os.path.join(train_img_dir, f"bright_{img_file}")
        bright_img.save(bright_img_path)
        
        # 亮度增强不改变标签位置
        bright_label_path = os.path.join(train_label_dir, f"bright_{img_file.rsplit('.', 1)[0]}.txt")
        with open(bright_label_path, 'w') as f:
            f.write(label_content)
        
        # 对比度增强
        enhancer = ImageEnhance.Contrast(img)
        contrast_img = enhancer.enhance(1.5)  # 增加对比度
        contrast_img_path = os.path.join(train_img_dir, f"contrast_{img_file}")
        contrast_img.save(contrast_img_path)
        
        # 对比度增强不改变标签位置
        contrast_label_path = os.path.join(train_label_dir, f"contrast_{img_file.rsplit('.', 1)[0]}.txt")
        with open(contrast_label_path, 'w') as f:
            f.write(label_content)

def main():
    """主函数"""
    print("开始准备水果识别数据集...")
    
    # 下载数据集
    download_dataset()
    
    # 创建YOLO格式的数据集
    create_yolo_dataset()
    
    # 数据增强
    augment_data()
    
    print("数据准备完成!")

if __name__ == "__main__":
    main() 